package org.usmile.algorithms.leetcode.middle;

import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * 973. 最接近原点的 K 个点
 *
 * 给定一个数组 points ，其中 points[i] = [xi, yi] 表示 X-Y 平面上的一个点，并且是一个整数 k ，返回离原点 (0,0) 最近的 k 个点。
 * 这里，平面上两点之间的距离是 欧几里德距离（ √(x1 - x2)2 + (y1 - y2)2 ）。
 * 你可以按 任何顺序 返回答案。除了点坐标的顺序之外，答案 确保 是 唯一 的。
 *
 * 示例 1：
 *
 * 输入：points = [[1,3],[-2,2]], k = 1
 * 输出：[[-2,2]]
 * 解释：
 * (1, 3) 和原点之间的距离为 sqrt(10)，
 * (-2, 2) 和原点之间的距离为 sqrt(8)，
 * 由于 sqrt(8) < sqrt(10)，(-2, 2) 离原点更近。
 * 我们只需要距离原点最近的 K = 1 个点，所以答案就是 [[-2,2]]。
 *
 * 示例 2：
 * 输入：points = [[3,3],[5,-1],[-2,4]], k = 2
 * 输出：[[3,3],[-2,4]]
 * （答案 [[-2,4],[3,3]] 也会被接受。）
 *
 * 提示：
 * 1 <= k <= points.length <= 104
 * -104 < xi, yi < 104
 */
public class _0973 {
}

class _0973_Solution {
    private final Random random = new Random(System.currentTimeMillis());
    // 前 k 个最大元素 --> 升序 --> 拿到后面 k 个元素即可
    // 前 k 个最小元素 --> 降序 --> 拿到后面 k 个元素即可
    public int[][] kClosest(int[][] points, int k) {
        int n = points.length;
        int target = n - k;

        int left = 0, right = n - 1;
        while (true) {
            int index = partition(points, left, right);
            if (index == target) {
                break;
            } else if (index < target) {
                left = index + 1;
            } else {
                right = index - 1;
            }
        }

        return Arrays.copyOfRange(points, n - k, n);
    }

    private int partition(int[][] points, int left, int right) {
        if (right > left) {
            int pivotIndex = left + 1 + random.nextInt(right - left);
            swap(points, pivotIndex, right);
        }
        int pivot = points[right][0] * points[right][0]
                + points[right][1] * points[right][1];
        int less = left, great = left;
        for (; great < right; great++) {
            int num = points[great][0] * points[great][0]
                    + points[great][1] * points[great][1];
            // 降序排列
            if (num >= pivot) {
                swap(points, less, great);
                less++;
            }
        }
        swap(points, less, right);

        return less;
    }

    private void swap(int[][] nums, int i, int j) {
        int[] tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

class _0973_Solution1 {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> pq
                = new PriorityQueue<>(k + 1, (a, b) -> (b[0] * b[0] + b[1] * b[1]) - (a[0] * a[0] + a[1] * a[1]));

        for (int[] point : points) {
            pq.add(point);
            if (pq.size() > k) pq.remove();
        }

        int[][] result = new int[k][2];
        int i = 0;
        while (!pq.isEmpty()) {
            result[i++] = pq.remove();
        }

        return result;
    }
}